import networkx as nx
import matplotlib.pyplot as plt


def main():
    G = nx.Graph()
    # topology construction logic
    G.add_node('1', type='host', info='10.25.131.11')
    G.add_node('2', type='host', info='10.25.131.12')
    G.add_node('3', type='host', info='10.25.131.13')

    G.add_node('11', type='alarm', info='10.25.131.11 LIVE_EPG 不通')
    G.add_node('22', type='alarm', info='10.25.131.11 LIVE_EPG 连不上')
    G.add_node('33', type='alarm', info='10.25.131.12 LIVE_EPG 连不上')
    G.add_node('44', type='alarm', info='10.25.131.12 VOD_EPG 不通')
    G.add_node('55', type='alarm', info='10.25.131.13 VOD_EPG 不通')

    G.add_node('111', type='business', info='LIVE_EPG')
    G.add_node('222', type='business', info='VOD_EPG')

    G.add_edge('1', '222', name='h_b')
    G.add_edge('2', '111', name='h_b')
    G.add_edge('3', '111', name='h_b')
    G.add_edge('1', '11', name='h_a')
    G.add_edge('1', '22', name='h_a')
    G.add_edge('2', '33', name='h_a')
    G.add_edge('2', '44', name='h_a')
    G.add_edge('3', '55', name='h_a')

    # G.add_edge('1', '222')
    # G.add_edge('2', '111')
    # G.add_edge('3', '111')
    # G.add_edge('1', '11')
    # G.add_edge('1', '22')
    # G.add_edge('2', '33')
    # G.add_edge('2', '44')
    # G.add_edge('3', '55')

    # G.add_edge('1', '222', name='h_b1', weight=1)
    # G.add_edge('2', '111', name='h_b2', weight=1)
    # G.add_edge('3', '111', name='h_b3', weight=1)
    # G.add_edge('1', '11', name='h_a1', weight=1)
    # G.add_edge('1', '22', name='h_a2', weight=1)
    # G.add_edge('2', '33', name='h_a3', weight=1)
    # G.add_edge('2', '44', name='h_a4', weight=1)
    # G.add_edge('3', '55', name='h_a5', weight=1)

    # draw graph with labels
    pos = nx.spring_layout(G)
    nx.draw(G, pos)

    # generate node_labels manually
    node_labels = {}
    for node in G.nodes:
        print(node, G.nodes[node])
        if G.nodes[node]["type"] == "host":
            node_labels[node] = 'host_id: ' + str(node)  # whatever you like
        if G.nodes[node]["type"] == "alarm":
            node_labels[node] = 'alarm_id: ' + str(node)  # whatever you like
        if G.nodes[node]["type"] == "business":
            node_labels[node] = 'business_id: ' + str(node)  # whatever you like

    nx.draw_networkx_labels(G, pos, labels=node_labels, font_color="red")
    nx.draw_networkx_edge_labels(G, pos,
                                 font_color="blue")  # no edge_labels parameter, default is showing all attributes of edges
    plt.show()


if __name__ == '__main__':
    main()
